import argparse

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='示例脚本：接收CUDA设备和端口参数')
    
    # 添加参数
    parser.add_argument(
        '-c',
        '--cuda_id', 
        type=str,
        required=True,  # 必须传入
        help='CUDA设备ID，例如 "0" 或 "0,1"（字符串类型）'
    )
    parser.add_argument(
        '-p',
        '--port', 
        type=int,
        required=True,  # 必须传入
        help='port 端口'
    )
    
    # 解析参数
    args = parser.parse_args()
    return args

args = parse_args()

import os,sys,pdb
os.environ['CUDA_VISIBLE_DEVICES']=args.cuda_id

from flask import Flask, request, send_file
import json
from util_flux import horizontal_concat_images
import numpy as np
from io import BytesIO
from PIL import Image

from demo_sam import load_sam2,get_mask_by_img_coor_inputbox
from flask_cors import CORS  # 导入 CORS

app = Flask(__name__)
CORS(app)  # 允许所有域名访问（开发环境推荐）

predictor = load_sam2()

img, control_img, rec_coor = None,None,None
save_dir = '/mnt/nas/shengjie/datasets/puthere_data_process'

# 这是你的图像分割函数，这里只是一个示例
def segment_image(image, rect):
    # 示例：简单地在矩形区域绘制一个红色矩形
    # 实际应用中替换为你的分割算法

    # 确保 image 是 PIL.Image.Image 类型，如果不是则转换
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    
    # 将图像转换为OpenCV格式
    # img_array = np.array(image)
    # img_cv = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
    
    # 绘制矩形
    x, y, w, h = int(rect['x']), int(rect['y']), int(rect['width']), int(rect['height'])
    # cv2.rectangle(img_cv, (x, y), (x + w, y + h), (0, 0, 255), 2)
    
    # 转换回PIL格式
    # result_img = Image.fromarray(cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB))

    mask = get_mask_by_img_coor_inputbox(predictor, image, x,y ,x+w,y+h)

    # res = Composite ( mask, image ) 

    return mask

@app.route('/process', methods=['POST'])
def process_image():
    if 'image' not in request.files or 'rect' not in request.form:
        return "Missing image or rectangle data", 400
    
    # 获取上传的图像
    file = request.files['image']
    image = Image.open(file.stream)
    
    # 获取矩形坐标
    rect = json.loads(request.form['rect'])

    print('rect::>>', rect)
    
    # 调用分割函数
    result_image = segment_image(image, rect) # np.unique = 0 1 
    # INSERT_YOUR_CODE
    # 确保输出为PIL.Image类型
    if isinstance(result_image, np.ndarray):
        # INSERT_YOUR_CODE
        # 将0/1 mask转为RGBA格式的PIL mask
        # 0->全透明，1->全不透明白色
        mask = (result_image > 0).astype(np.uint8) * 255  # 0/1 -> 0/255
        # INSERT_YOUR_CODE
        # 把mask存入tmp_mask.png中查看效果
        # Image.fromarray(mask).save("tmp_mask.png")
        rgba = np.zeros((mask.shape[0], mask.shape[1], 4), dtype=np.uint8)
        # RGB三个通道用原图代替, 最终的图就是被遮罩的原图
        rgba[..., 0:3] = np.array(image.convert("RGB"))  # 用原图RGB通道
        rgba[..., 3] = mask   # A通道为mask
        result_image = Image.fromarray(rgba, mode='RGBA')
    
    # 将结果图像转换为字节流返回
    img_io = BytesIO()
    result_image.save(img_io, 'PNG')
    img_io.seek(0)

    global img, control_img, rec_coor
    img = image.convert('RGB')
    # control_img = mask==255的原图部分
    # 先获取mask（A通道），再用mask把原图抠出来
    if result_image.mode == 'RGBA':
        rgba_arr = np.array(result_image)
        mask = rgba_arr[..., 3]  # A通道
        rgb = np.array(image.convert("RGB"))
        control_img_arr = np.ones_like(rgb)*255
        control_img_arr[mask == 255] = rgb[mask == 255]
        control_img = Image.fromarray(control_img_arr, mode='RGB')
    else:
        control_img = result_image.convert('RGB')
    rec_coor = rect
    horizontal_concat_images([img, control_img]).save('tmp_put.jpg')
    
    return send_file(img_io, mimetype='image/png')

# 保存三个信息
# 1 矩形框坐标 rec coor
# 2 ori img
# 3 res img
# 处理过程 
# cropped = crop(res img, rec coor)
# control img = paste(ori img, cropped)
# concat(ori img, control img).save(save_dir, $time.png)
# INSERT_YOUR_CODE
@app.route('/save', methods=['POST'])
def save_images():
    import os
    from datetime import datetime

    # 假设全局变量 img, control_img, rec_coor 已经在 /process 中被赋值
    global img, control_img, rec_coor,save_dir

    # 检查变量是否存在
    if img is None or control_img is None or rec_coor is None:
        return "No image data to save. Please process an image first.", 400

    # 1. 裁剪 control_img 的 rec_coor 区域
    # rec_coor 应为 dict: {x, y, width, height}
    x = int(rec_coor.get('x', 0))
    y = int(rec_coor.get('y', 0))
    w = int(rec_coor.get('width', 0))
    h = int(rec_coor.get('height', 0))
    box = (x, y, x + w, y + h)
    cropped = control_img.crop(box)

    # 2. 将 cropped 区域粘贴回 ori img 得到 control_img_pasted
    control_img_pasted = img.copy()
    # 对 cropped 做随机仿射变换（旋转、缩放、平移等）用于数据增强，最后保持和输入的大小一致
    import random

    def random_affine_keep_size(img):
        # 输入 img: PIL.Image，输出和输入同样大小的PIL.Image
        w, h = img.size
        # 随机旋转 -15~15度
        angle = random.uniform(-15, 15)
        # 随机缩放 0.9~1.1
        scale = random.uniform(0.9, 1.1)
        # 随机平移 -10~10像素
        max_dx = min(10, w // 10)
        max_dy = min(10, h // 10)
        dx = random.randint(-max_dx, max_dx)
        dy = random.randint(-max_dy, max_dy)
        # 构造仿射矩阵
        # PIL affine: (a, b, c, d, e, f) -> x' = a*x + b*y + c, y' = d*x + e*y + f
        # 先缩放+旋转
        import math
        angle_rad = math.radians(angle)
        a = scale * math.cos(angle_rad)
        b = scale * math.sin(angle_rad)
        d = -scale * math.sin(angle_rad)
        e = scale * math.cos(angle_rad)
        # 平移
        c = dx
        f = dy
        # 中心变换
        cx, cy = w/2, h/2
        # 为了绕中心旋转缩放，需先平移到原点，仿射后再平移回去
        # 先平移到原点
        # 复合仿射矩阵的平移分量
        c = -a*cx - b*cy + cx + dx
        f = -d*cx - e*cy + cy + dy
        affine_matrix = (a, b, c, d, e, f)
        img_affine = img.transform((w, h), Image.AFFINE, affine_matrix, resample=Image.BICUBIC, fillcolor=(255,255,255))
        return img_affine

    cropped_aug = random_affine_keep_size(cropped)
    control_img_pasted.paste(cropped_aug, box)

    # 3. 拼接 ori img 和 control_img_pasted
    # 统一为RGB
    ori_img_rgb = img.convert("RGB")
    control_img_rgb = control_img_pasted.convert("RGB")
    w1, h1 = ori_img_rgb.size
    w2, h2 = control_img_rgb.size
    concat_img = Image.new("RGB", (w1 + w2, max(h1, h2)))
    concat_img.paste(ori_img_rgb, (0, 0))
    concat_img.paste(control_img_rgb, (w1, 0))

    # 4. 保存到 save_dir, 文件名为时间戳.jpg
    os.makedirs(save_dir, exist_ok=True)
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    save_path = os.path.join(save_dir, f"{timestamp}.jpg")
    concat_img.save(save_path)
    # INSERT_YOUR_CODE
    # 将 box 信息保存到与图片同名的 txt 文件中
    box_txt_path = os.path.splitext(save_path)[0] + ".txt"
    with open(box_txt_path, "w") as f:
        f.write(str(box))
    concat_img.save('tmp_put.jpg')

    return f"Saved to {save_path}", 200


if __name__ == '__main__':
    app.run(debug=False,host='0.0.0.0',port=args.port)

